#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/8/2 3:23 下午
# @Author  : WangZhixing
import torch
import torch.nn.functional as F
from torch.nn import Sequential as Seq, Linear as Lin, ReLU, BatchNorm1d as BN, Dropout

from Model.Module.MLP import MLP


class Discriminator1(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(Discriminator1, self).__init__()
        self.lin1 = torch.nn.Linear(in_channels, hidden_channels)
        self.lin2 = torch.nn.Linear(hidden_channels, hidden_channels)
        self.lin3 = torch.nn.Linear(hidden_channels, out_channels)

    def forward(self, x):
        x = F.relu(self.lin1(x))
        x = F.relu(self.lin2(x))
        x = self.lin3(x)
        return x

class Discriminator2(torch.nn.Module):
    def __init__(self, in_channels, out_channels=1):
        super(Discriminator2, self).__init__()
        self.mlp = Seq(
            MLP([in_channels, 128, 64]), Dropout(0.5),
            Lin(64, out_channels))

    def forward(self, x):
        return self.mlp(x)